import re
import unittest

import numpy as np

import MatPlot
import ProductEval
from Asm10Model import Asm10Model
from AsmExcel import AsmExcel
from Fy3dL1Hdf import Fy3dL1Hdf, Fy3dL1HdfBatch
from Fy3dL2VsmHdf import Fy3dL2VsmHdf, Fy3dL2VsmHdfBatch
from Fy3dLd3 import Fy3dLd3, Fy3dLd3Batch
from GenericExcel import GenericExcel
from GeoTiff import GeoTiffBatch, GeoTiff
from NvdiLd2 import NvdiLd2
from Timer import Timer
from utils import stat_data, find_files, last_days, find_last_files

ASM_DIR = "data/ASM_HN"
LD3_DIR = 'data/FY3D MERSI/LD3'
LD3_OK_DIR = 'data/FY3D MERSI/LD3/ok'
LD3_ERR_DIR = 'data/FY3D MERSI/LD3/error'
NDVI_DIR = 'data/FY3D MERSI/NDVI'
LST_DIR = 'data/FY3D MERSI/LST'
TVDI_DIR = 'data/FY3D MERSI/TVDI'


class AllTests(unittest.TestCase):

    def test_1_1(self):
        ld3 = Fy3dLd3("data/FY3D MERSI/LD3/FY3D_X_2022_01_01_14_24_A_G_MERSI_1000M_L1B_proj_Latlon.ld3")
        ld3.generate_three(NDVI_DIR, LST_DIR, TVDI_DIR)
        pass

    def test_1_2(self):
        ld3_batch = Fy3dLd3Batch("data/FY3D MERSI/LD3")
        ld3_batch.generate_three(NDVI_DIR, LST_DIR, TVDI_DIR)
        pass

    def test_1_2_(self):
        # 只处理当前目录下文件，成功移到ok目录，失败移到ERROR目录，子目录下的文件不处理
        ld3_batch = Fy3dLd3Batch("data/FY3D MERSI/LD3")
        ld3_batch.generate_three(NDVI_DIR, LST_DIR, TVDI_DIR, LD3_OK_DIR, LD3_ERR_DIR)
        pass

    def test_1_3(self):
        # NDVI值域 [-1,1]，添加最大值合成、平均值合成功能，
        ndvi_tiffs = GeoTiffBatch("data/FY3D MERSI/NDVI")
        max_value = ndvi_tiffs.max_at_mn(100, 100)
        mean_value = ndvi_tiffs.mean_at_mn(100, 100)
        print(max_value, mean_value)
        pass

    def test_1_4(self):
        l1 = Fy3dL1Hdf("data/FY3D_MWRI_L1/ASCEND/20220101/FY3D_MWRIA_GBAL_L1_20220101_0045_010KM_MS.HDF")
        l1.print()
        pass

    def test_1_5(self):
        timer = Timer()
        timer.tick("read hdf file")
        l1 = Fy3dL1Hdf("data/FY3D_MWRI_L1/ASCEND/20220101/FY3D_MWRIA_GBAL_L1_20220101_0045_010KM_MS.HDF")
        timer.tick("start map_to_area")
        l1.map_to_area()
        timer.tick("end map_to_area")
        x, y = l1.latlon_to_xy(34.4919, 113.0503)
        result = l1.earth_ob_at_xy(x, y)
        print(x, y, result)
        pass

    def test_1_6(self):
        l1_batch = Fy3dL1HdfBatch("data/FY3D_MWRI_L1/ASCEND/20220101")
        l1_batch.map_to_area()
        x, y = l1_batch.latlon_to_xy(34.4919, 113.0503)
        print(x, y)
        result = l1_batch.earth_ob_at_xy(x, y)
        print(result)
        pass

    def test_1_7(self):
        l2_vsm = Fy3dL2VsmHdf("data/FY3D_MWRI_VSM/FY3D_MWRIX_GBAL_L2_VSM_MLT_ESD_20220101_POAD_025KM_MS.HDF")
        stat_data(l2_vsm.vsm_a_data, np.nan, " VSM_A ")
        stat_data(l2_vsm.vsm_d_data, np.nan, " VSM_D ")
        print()
        x, y = l2_vsm.latlon_to_xy(34.4919, 113.0503)
        print(x, y)
        vsm_a_value = l2_vsm.vsm_a_at(34.4919, 113.0503)
        print(vsm_a_value, l2_vsm.vsm_a_data[int(x), int(y)])
        vsm_d_value = l2_vsm.vsm_d_at(34.4919, 113.0503)
        print(vsm_d_value, l2_vsm.vsm_d_data[int(x), int(y)])
        pass

    def test_1_8(self):
        l2_vsm_batch = Fy3dL2VsmHdfBatch("data/FY3D_MWRI_VSM")
        mean_value = l2_vsm_batch.mean_vsm_a_at_mn(0, 0)
        max_value = l2_vsm_batch.max_vsm_a_at_mn(0, 0)
        print(mean_value, max_value)
        pass

    def test_1_9(self):
        nvdi_ld2 = NvdiLd2("data/NDVI_DEK/NDVI_20220110_max.ld2")
        pass

    def test_1_10(self):
        tiff = GeoTiff("data/landuse2020/lucc2020_Dxx_hn.tif")
        value = tiff.get_pixel_value(34.4919, 113.0503)
        print(value)
        pass

    def test_1_11(self):
        tiff = GeoTiff("data/soiltype/soil_hn.tif")
        value = tiff.get_pixel_value(34.4919, 113.0503)
        print(value)
        pass

    def test_1_12(self):
        asm_xls = AsmExcel("data/ASM_HN/ASM_202201_10.xls")
        print(asm_xls.df)
        pass

    def test_1_13(self):
        tvdi_tiff = GeoTiff("data/FY3D MERSI/TVDI/FY3D_X_2022_01_01_14_24_A_G_MERSI_1000M_L1B_proj_Latlon_TVDI.TIF")
        value = tvdi_tiff.get_pixel_value_interpolated(34.4919, 113.0503)
        print(value)
        pass

    def test_1_14(self):
        tvdi_tiff = GeoTiff("data/FY3D MERSI/TVDI/FY3D_X_2022_01_01_14_24_A_G_MERSI_1000M_L1B_proj_Latlon_TVDI.TIF")
        MatPlot.show_data(tvdi_tiff.band1_data)
        MatPlot.export_png(tvdi_tiff.band1_data,
                           "temp/FY3D_X_2022_01_01_14_24_A_G_MERSI_1000M_L1B_proj_Latlon_TVDI.png")
        height, width = tvdi_tiff.band1_data.shape
        MatPlot.export_png(tvdi_tiff.band1_data,
                           "temp/FY3D_X_2022_01_01_14_24_A_G_MERSI_1000M_L1B_proj_Latlon_TVDI_1.png", width, height,
                           True)
        pass

    def test_1_15(self):
        geo_tiff = GeoTiff("data/FY3D MERSI/TVDI/FY3D_X_2022_01_01_14_24_A_G_MERSI_1000M_L1B_proj_Latlon_TVDI.TIF")
        MatPlot.show_geo_data(geo_tiff.band1_data, geo_tiff.geo_transform, geo_tiff.projection,
                              "data/soiltype/hn_sj.shp")
        MatPlot.export_geo_png(geo_tiff.band1_data, geo_tiff.geo_transform, geo_tiff.projection,
                               "temp/FY3D_X_2022_01_01_14_24_A_G_MERSI_1000M_L1B_proj_Latlon_TVDI.png")
        height, width = geo_tiff.band1_data.shape
        MatPlot.export_geo_png(geo_tiff.band1_data, geo_tiff.geo_transform, geo_tiff.projection,
                               "temp/FY3D_X_2022_01_01_14_24_A_G_MERSI_1000M_L1B_proj_Latlon_TVDI_1.png",
                               width, height, True, "data/soiltype/hn_sj.shp")
        pass

    def test_1_15_(self):
        # TVDI保存成左图样子
        geo_tiff = GeoTiff("data/FY3D MERSI/TVDI/FY3D_X_2022_01_01_14_24_A_G_MERSI_1000M_L1B_proj_Latlon_TVDI.TIF")
        height, width = geo_tiff.band1_data.shape
        MatPlot.export_geo_png(geo_tiff.band1_data, geo_tiff.geo_transform, geo_tiff.projection,
                               "temp/FY3D_X_2022_01_01_14_24_A_G_MERSI_1000M_L1B_proj_Latlon_TVDI_2.png",
                               width, height, False, "data/soiltype/hn_sj.shp")
        pass

    def test_1_16(self):
        # VSM也加个图像输出
        l2_vsm = Fy3dL2VsmHdf("data/FY3D_MWRI_VSM/FY3D_MWRIX_GBAL_L2_VSM_MLT_ESD_20220101_POAD_025KM_MS.HDF")
        height, width = l2_vsm.vsm_a_data.shape
        MatPlot.export_geo_png(l2_vsm.vsm_a_data, l2_vsm.geo_transform, l2_vsm.projection,
                               "temp/FY3D_MWRIX_GBAL_L2_VSM_MLT_ESD_20220101_POAD_025KM_MS_A.png")
        MatPlot.export_geo_png(l2_vsm.vsm_a_data, l2_vsm.geo_transform, l2_vsm.projection,
                               "temp/FY3D_MWRIX_GBAL_L2_VSM_MLT_ESD_20220101_POAD_025KM_MS_A_ha.png",
                               width, height, False, "data/soiltype/hn_sj.shp")
        MatPlot.export_geo_png(l2_vsm.vsm_d_data, l2_vsm.geo_transform, l2_vsm.projection,
                               "temp/FY3D_MWRIX_GBAL_L2_VSM_MLT_ESD_20220101_POAD_025KM_MS_D.png")
        MatPlot.export_geo_png(l2_vsm.vsm_d_data, l2_vsm.geo_transform, l2_vsm.projection,
                               "temp/FY3D_MWRIX_GBAL_L2_VSM_MLT_ESD_20220101_POAD_025KM_MS_D_ha.png",
                               width, height, False, "data/soiltype/hn_sj.shp")
        pass

    def test_2_1(self):
        df = ProductEval.combine_asm10_and_tvdi("data/ASM_HN/ASM_202201_10.xls", "data/FY3D MERSI/TVDI",
                                                "temp/202201_TVDI_ASM_10.xlsx")
        print(df)
        pass

    def test_2_2(self):
        df = ProductEval.combine_asm10_and_fy3dvsm("data/ASM_HN/ASM_202201_10.xls", "data/FY3D_MWRI_VSM",
                                                   "temp/202201_VSM_ASM_10.xlsx")
        print(df)
        df = ProductEval.combine_asm10_and_fy3dvsm("data/ASM_HN/ASM_202202_10.xls", "data/FY3D_MWRI_VSM",
                                                   "temp/202202_VSM_ASM_10.xlsx")
        print(df)
        pass

    def test_2_3(self):
        # ASM已经缩小100倍
        df = ProductEval.clean_asm10_vsm("temp/202201_VSM_ASM_10.xlsx", "temp/cleaned_202201_VSM_ASM_10.xlsx")
        print(df)
        df = ProductEval.clean_asm10_vsm("temp/202202_VSM_ASM_10.xlsx", "temp/cleaned_202202_VSM_ASM_10.xlsx")
        print(df)
        # 清洗掉的数据是什么？
        #   asm10_vsm.drop_rows(lambda r: isna(r['FY3D_A']) or isna(r['FY3D_D']))
        #   FY3D_A 或 FY3D_D 有一个是无效值则被清洗
        pass

    def test_2_4(self):
        asm10_vsm_filename_pattern = r"cleaned_\d{6}_VSM_ASM_10\.xlsx"
        asm10_vsm_files = find_files("temp", lambda f: bool(re.search(asm10_vsm_filename_pattern, f)))
        df = ProductEval.concat_asm10_vsm_files(asm10_vsm_files, "temp/cleaned_com_VSM_ASM_10.xlsx")
        print(df)
        pass

    def test_2_5(self):
        result_df = ProductEval.analyse_asm10_vsm_com("temp/cleaned_com_VSM_ASM_10.xlsx",
                                                      "temp/cleaned_com_VSM_ASM_10_analysis.xlsx")
        print(result_df)
        pass

    def test_3_1_prepare_train_data(self):
        model1 = Asm10Model("model/model1")
        model1.prepare_train_data(ASM_DIR,
                                  tvdi_dir=TVDI_DIR,
                                  ndvi_dir=NDVI_DIR,
                                  soiltype_file="data/soiltype/soil_hn.tif",
                                  landuse_file="data/landuse2020/lucc2020_Dxx_hn.tif",
                                  save_path="temp/model1_train_data.xlsx")
        pass

    def test_3_1_fit(self):
        model1 = Asm10Model("model/model1")
        input_xls = GenericExcel("temp/model1_train_data.xlsx", {'soiltype': 'category', 'landuse': 'category'})
        leaderboards = model1.fit(input_xls.df,
                                  ['tvdi', 'ndvi', 'soiltype', 'landuse'],
                                  # 注意结果列名中不可以出现中文 ["VWC", "SRH", "GWC", "AWS"]
                                  ["VWC", "SRH"])
        print(leaderboards)
        pass

    def test_3_1_prepare_data(self):
        model1 = Asm10Model("model/model1")
        model1.prepare_data('2022_03_23_12',
                            tvdi_dir=TVDI_DIR,
                            ndvi_dir=NDVI_DIR,
                            soiltype_file="data/soiltype/soil_hn.tif",
                            landuse_file="data/landuse2020/lucc2020_Dxx_hn.tif",
                            save_path="temp/2022_03_23_12_data.xlsx")
        pass

    def test_3_1_predict(self):
        model1 = Asm10Model("model/model1")
        input_xls = GenericExcel("temp/2022_03_23_12_data.xlsx", {'soiltype': 'category', 'landuse': 'category'})
        df = model1.predict(input_xls.df,
                            ['tvdi', 'ndvi', 'soiltype', 'landuse'],
                            save_path="temp/2022_03_23_12_result.xlsx")
        print(df)
        pass

    def test_3_1_predict_as_tiffs(self):
        model1 = Asm10Model("model/model1")
        input_data = model1.prepare_data('2022_03_23_12',
                                         tvdi_dir=TVDI_DIR,
                                         ndvi_dir=NDVI_DIR,
                                         soiltype_file="data/soiltype/soil_hn.tif",
                                         landuse_file="data/landuse2020/lucc2020_Dxx_hn.tif",
                                         save_path="temp/2022_03_23_12_data.xlsx")
        result_df = model1.predict(input_data,
                                   ['tvdi', 'ndvi', 'soiltype', 'landuse'],
                                   save_path="temp/2022_03_23_12_result.xlsx")
        model1.export_tiffs(result_df,
                            sample_tiff=None,  # 如果之前进行了prepare_data，则可以自动提取样例tiff文件，否则需要指定路径
                            save_path_func=lambda label: f"temp/2022_03_23_12_{label}.tiff")
        pass

    def test_3_2_prepare(self):
        model2 = Asm10Model("model/model1")
        model2.prepare_train_data(ASM_DIR,
                                  l1_dir="data/FY3D_MWRI_L1",
                                  ndvi_dir=NDVI_DIR,
                                  lst_dir=LST_DIR,
                                  soiltype_file="data/soiltype/soil_hn.tif",
                                  landuse_file="data/landuse2020/lucc2020_Dxx_hn.tif",
                                  save_path="temp/model2_train_data.xlsx")
        pass

    def test_3_2_fit(self):
        model2 = Asm10Model("model/model2")
        input_xls = GenericExcel("temp/model2_train_data.xlsx", {'soiltype': 'category', 'landuse': 'category'})
        model2.fit(input_xls.df,
                   ['earth_ob_1', 'earth_ob_2', 'earth_ob_3', 'earth_ob_4', 'earth_ob_5',
                    'earth_ob_6', 'earth_ob_7', 'earth_ob_8', 'earth_ob_9', 'earth_ob_10',
                    'ndvi', 'lst', 'soiltype', 'landuse'],
                   # 注意结果列名中不可以出现中文 ["VWC", "SRH", "GWC", "AWS"]
                   ["VWC"])
        pass

    def test_3_2_predict(self):
        model2 = Asm10Model("model/model2")
        input_xls = GenericExcel("temp/model2_train_data.xlsx", {'soiltype': 'category', 'landuse': 'category'})
        df = model2.predict(input_xls.df,
                            ['earth_ob_1', 'earth_ob_2', 'earth_ob_3', 'earth_ob_4', 'earth_ob_5',
                             'earth_ob_6', 'earth_ob_7', 'earth_ob_8', 'earth_ob_9', 'earth_ob_10',
                             'ndvi', 'lst', 'soiltype', 'landuse'],
                            save_path="temp/model2_result_data.xlsx")
        print(df)
        pass

    def test_3_4_prepare(self):
        model4 = Asm10Model("model/model4")
        model4.prepare_train_data(ASM_DIR,
                                  vsm_dir="data/FY3D_MWRI_VSM",
                                  ndvi_dir=NDVI_DIR,
                                  soiltype_file="data/soiltype/soil_hn.tif",
                                  landuse_file="data/landuse2020/lucc2020_Dxx_hn.tif",
                                  save_path="temp/model4_train_data.xlsx")
        pass

    def test_3_4_fit(self):
        model4 = Asm10Model("model/model4")
        input_xls = GenericExcel("temp/model4_train_data.xlsx", {'soiltype': 'category', 'landuse': 'category'})
        model4.fit(input_xls.df,
                   ['ndvi', 'soiltype', 'landuse', "VWC"],
                   ['vsm'])
        pass

    def test_3_4_predict(self):
        model4 = Asm10Model("model/model4")
        input_xls = GenericExcel("temp/model4_train_data.xlsx", {'soiltype': 'category', 'landuse': 'category'})
        df = model4.predict(input_xls.df,
                            ['ndvi', 'soiltype', 'landuse', "VWC"],
                            save_path="temp/model4_result_data.xlsx")
        print(df)
        pass

    def test_hdf_vs_tiff_0(self):
        hdf = Fy3dL2VsmHdf("data/FY3D_MWRIX_GBAL_L2_VSM_MLT_ESD_20240102_POAD_025KM_MS.HDF")
        tiff_a = GeoTiff("data/FY3D_MWRI_A_20240102.TIFF")
        MatPlot.show_data(hdf.vsm_a_data)
        MatPlot.show_data(tiff_a.band1_data)
        tiff_d = GeoTiff("data/FY3D_MWRI_D_20240102.TIFF")
        MatPlot.show_data(hdf.vsm_d_data)
        MatPlot.show_data(tiff_d.band1_data)
        pass

    def test_hdf_vs_tiff_1(self):
        hdf = Fy3dL2VsmHdf("data/FY3D_MWRIX_GBAL_L2_VSM_MLT_ESD_20240102_POAD_025KM_MS.HDF")
        stat_data(hdf.vsm_a_data, np.nan, " VSM_A ")
        lat = 34.4919
        lon = 113.0503
        print("(", lat, ",", lon, ")\n")
        a0 = hdf.vsm_a_at(lat, lon)
        d0 = hdf.vsm_d_at(lat, lon)
        tiff_a = GeoTiff("data/FY3D_MWRI_A_20240102.TIFF")
        a1 = tiff_a.get_pixel_value_interpolated(lat, lon)
        tiff_d = GeoTiff("data/FY3D_MWRI_D_20240102.TIFF")
        d1 = tiff_d.get_pixel_value_interpolated(lat, lon)
        print("(", a0, ",", d0, ") (", a1, ",", d1, ")")
        pass

    def test_tiff_batch(self):
        # 获取最近10天的文件路径
        file_paths = find_last_files("data/FY3D MERSI/LST",
                                     lambda f: f.upper().endswith('.TIF'),
                                     10, '2022_02_26')
        # 创建GeoTiffBatch对象
        tiffs = GeoTiffBatch(tiff_paths=file_paths)

        # 定义有效值检查函数
        check_func = lambda v: v != -98 and v != -99
        # 执行合并操作
        tiffs.do_overlap(check_func, save_path="data/FY3D MERSI_COM/LST/2022_02_26_overlapped.tiff")
        # 执行聚合操作 - 最大值
        tiffs.do_aggregate(check_func, aggregate_func=np.nanmax,
                           save_path="data/FY3D MERSI_COM/LST/2022_02_26_max.tiff")
        # # 执行聚合操作 - 平均值
        # tiffs.do_aggregate(check_func, aggregate_func=np.nanmean,
        #                    save_path="data/FY3D MERSI_COM/LST/2022_02_26_mean.tiff")
        # # 执行聚合操作 - 中位数
        # tiffs.do_aggregate(check_func, aggregate_func=np.nanmedian,
        #                    save_path="data/FY3D MERSI_COM/LST/2022_02_26_median.tiff")
        tiffs.close()
        pass
